from typing import List, Union
import requests

# 请求头（无需 Authorization）
headers = {
    'Content-Type': 'application/json'
}

# 文献库检索接口说明
BASE_URL = "http://180.184.65.98:38880/atomgit"

# API列表
# 序号  功能描述                                   路由                          输入参数                                      返回值
# 1      根据查询文本搜索论文片段                   /search_papers                query: str (必填), top_k: int = 30 (可选)     JSON 格式：包含匹配的论文片段信息
# 2      根据论文唯一 ID 获取论文片段               /query_by_paper_id            paper_id: str (必填), top_k: int = 5 (可选)   JSON 格式：包含匹配的论文片段信息
# 3      根据论文标题精确匹配获取论文片段           /query_by_title               title: str (必填), top_k: int = 100 (可选)    JSON 格式：包含匹配的论文片段信息
# 4      获取论文数据库的元数据信息                 /metadata                     无                                           JSON 格式：元数据信息
# 5      根据论文标题中包含的文本查询论文片段       /query_by_title_contain       title: str (必填), top_k: int = 1000 (可选)   JSON 格式：包含匹配的论文片段信息
# 6      根据论文内容片段中包含的文本查询论文片段   /query_by_chunk               chunk: str (必填), top_k: int = 50 (可选)     JSON 格式：包含匹配的论文片段信息
# 7      根据相似标题查找论文                       /query_by_title_like        title: str (必填), top_k: int = 5 (可选)     JSON 格式：包含相似论文信息列表 
# 8      根据关键词查找论文ID和标题                 /query_by_keyword           keyword (str):（必填）                        JSON 格式：包含论文ID和标题列表
# 9      根据论文ID获取完整文本                     /query_whole_text_by_id     paper_id: str (必填)                          论文完整文本 (字符串) 或 None
# 10     根据论文标题获取完整文本                   /query_whole_text_by_title  title: str (必填)                             论文完整文本 (字符串) 或 None
# 11     根据论文ID获取关键词列表                   /query_keywords_by_id       paper_id: str (必填)                          JSON格式：关键词列表 (字符串列表) 或 None
# 12     根据论文标题获取关键词列表                 /query_keywords_by_title    title: str (必填)                               JSON格式：关键词列表 (字符串列表) 或 None
# 13     查看关键词统计信息                         /keywords_metadata          无                                             JSON格式：关键词列表 (字符串列表) 或 None

def search_papers(query: str, top_k: int = 30) -> dict:
    """
    根据查询文本搜索论文片段。

    参数：
        query (str): 查询文本（必填）。
        top_k (int): 返回的论文片段数量，默认值为 30（可选）。

    返回：
        JSON 格式：包含匹配的论文片段信息。
    """
    url = f"{BASE_URL}/search_papers"
    data = {"query": query, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def get_paper_snippets_by_id(paper_id: str, top_k: int = 5) -> dict:
    """
    根据论文唯一 ID 获取论文片段。

    参数：
        paper_id (str): 论文唯一 ID（必填）。
        top_k (int): 返回的论文片段数量，默认值为 5（可选）。

    返回：
        JSON 格式：包含匹配的论文片段信息。
    """
    url = f"{BASE_URL}/query_by_paper_id"
    data = {"paper_id": paper_id, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def get_paper_snippets_by_title(title: str, top_k: int = 100) -> dict:
    """
    根据论文标题精确匹配获取论文片段。

    参数：
        title (str): 论文标题（必填）。
        top_k (int): 返回的论文片段数量，默认值为 100（可选）。

    返回：
        JSON 格式：包含匹配的论文片段信息。
    """
    url = f"{BASE_URL}/query_by_title"
    data = {"title": title, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def get_metadata() -> dict:
    """
    获取论文数据库的元数据信息。

    参数：
        无。

    返回：
        JSON 格式：元数据信息。
    """
    url = f"{BASE_URL}/metadata"

    rsp = requests.get(url, headers=headers)
    return rsp.json()


def search_papers_by_title_contain(title: str, top_k: int = 1000) -> dict:
    """
    根据论文标题中包含的文本查询论文片段。

    参数：
        title (str): 标题中包含的文本（必填）。
        top_k (int): 返回的论文片段数量，默认值为 1000（可选）。

    返回：
        JSON 格式：包含匹配的论文片段信息。
    """
    url = f"{BASE_URL}/query_by_title_contain"
    data = {"title": title, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def search_papers_by_chunk(chunk: str, top_k: int = 50) -> dict:
    """
    根据论文内容片段中包含的文本查询论文片段。

    参数：
        chunk (str): 文本片段（必填）。
        top_k (int): 返回的论文片段数量，默认值为 50（可选）。

    返回：
        JSON 格式：包含匹配的论文片段信息。
    """
    url = f"{BASE_URL}/query_by_chunk_contain"
    data = {"chunk": chunk, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def query_by_title_like(title: str, top_k: int = 5) -> dict:
    """
    根据相似标题查找论文。

    参数：
        title (str): 论文标题（必填）。
        top_k (int): 返回的相似论文数量，默认值为 5（可选）。

    返回：
        JSON 格式：包含相似论文信息列表。
    """
    url = f"{BASE_URL}/query_by_title_like"
    data = {"title": title, "top_k": top_k}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def query_by_keyword(keyword: str) -> dict:
    """
    根据关键词查找论文ID和标题。

    参数：
        keyword (str): 关键词（必填）。

    返回：
        JSON 格式：包含论文ID和标题列表。
    """
    url = f"{BASE_URL}/query_by_keyword"
    data = {"keyword": keyword}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json()


def query_whole_text_by_id(paper_id: str) -> str:
    """
    根据论文ID获取完整文本。

    参数：
        paper_id (str): 论文唯一 ID（必填）。

    返回：
        论文完整文本 (字符串) 或 None。
    """
    url = f"{BASE_URL}/query_whole_text_by_id"
    data = {"paper_id": paper_id}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.text if rsp.status_code == 200 else None


def query_whole_text_by_title(title: str) -> str:
    """
    根据论文标题获取完整文本。

    参数：
        title (str): 论文标题（必填）。

    返回：
        论文完整文本 (字符串) 或 None。
    """
    url = f"{BASE_URL}/query_whole_text_by_title"
    data = {"title": title}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.text if rsp.status_code == 200 else None


def query_keywords_by_id(paper_id: str) -> list:
    """
    根据论文ID获取关键词列表。

    参数：
        paper_id (str): 论文唯一 ID（必填）。

    返回：
        关键词列表 (字符串列表) 或 None。
    """
    url = f"{BASE_URL}/query_keywords_by_id"
    data = {"paper_id": paper_id}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json() if rsp.status_code == 200 else None


def query_keywords_by_title(title: str) -> list:
    """
    根据论文标题获取关键词列表。

    参数：
        title (str): 论文标题（必填）。

    返回：
        关键词列表 (字符串列表) 或 None。
    """
    url = f"{BASE_URL}/query_keywords_by_title"
    data = {"title": title}

    rsp = requests.get(url, params=data, headers=headers, proxies={"http": None, "https": None})
    return rsp.json() if rsp.status_code == 200 else None


def keywords_metadata() -> dict:
   """
   查看关键词统计信息。

   参数：无

   返回：
       JSON 格式：包含关键词及其计数的字典。
   """
   url = f"{BASE_URL}/keywords_metadata"

   rsp = requests.get(url, headers=headers)
   
   return rsp.json()

if __name__ == "__main__":
    import json
    from api_prompt_engineering import PromptEngineering
    PE = PromptEngineering()
    
    # 示例调用
    # print(PE.print_json(get_metadata())) 
    #print(search_papers("深度学习")) 
    # print(PE.print_json(get_paper_snippets_by_id("6294359f5aee126c0f2fe2b4")))   #paper-id有什么用 ?
    # get_paper_snippets_by_title("Deep Learning")
    # print(json.dumps(get_metadata(), ensure_ascii=False, indent=4))  
    # chunk=search_papers_by_title_contain("Text2SQL")
    # paper_list=[]
    # chunk=search_papers("Deep Learning Survey")
    # paper_ids = [item['entity']['paper_id'] for item in chunk if 'entity' in item and 'paper_id' in item['entity']]
    # for id in paper_ids:
    #     paper_list.append(get_paper_snippets_by_id(id))
    # with open("__caches/paper_list.json","w",encoding="utf-8")as f:
    #     f.write(json.dumps(paper_list, ensure_ascii=False, indent=4))

    # while True:
    #     print(" >>输入关键词检索: ", end="")
    #     search_keyword = input()
    #     chunk=search_papers(search_keyword)
    #     print(json.dumps(chunk, indent=4, ensure_ascii=False)[:4096]+"...")
    #     print("-"*100)